﻿using NgpSharp.RayMaching;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch;

namespace NgpSharp
{
    public class RaysCoordsFunction : autograd.MultiTensorFunction<RaysCoordsFunction>
    {
        /// <summary>
        /// 
        /// </summary>
        public override string Name => nameof(RaysCoordsFunction);

        public override List<Tensor> forward(autograd.AutogradContext ctx, params object[] vars)
        {
            using var c = torch.NewDisposeScope();

            Tensor rays_o = (Tensor)vars[0];
            Tensor rays_d = (Tensor)vars[1];
            Tensor aabb = (Tensor)vars[2];
            Tensor near = (Tensor)vars[3];
            Tensor far = (Tensor)vars[4];
            Tensor xyzs = (Tensor)vars[5];
            Tensor z_vals = (Tensor)vars[6];
            Tensor deltas = (Tensor)vars[7];
            float min_near = (float)vars[8];
            bool perturb = (bool)vars[9];

            uint num_steps = (uint)deltas.shape[^1];
            uint N = (uint)rays_o.shape[0];

            RayMachingLoader.rays_coords(rays_o.Handle, rays_d.Handle, aabb.Handle,
                num_steps, min_near, N, perturb,
                near.Handle, far.Handle, xyzs.Handle, z_vals.Handle, deltas.Handle);

            return
            [
                near.MoveToOuterDisposeScope(),
                far.MoveToOuterDisposeScope(),
                xyzs.MoveToOuterDisposeScope(),
                z_vals.MoveToOuterDisposeScope(),
                deltas.MoveToOuterDisposeScope()
            ];
        }

        public override List<Tensor> backward(autograd.AutogradContext ctx, List<Tensor> grad_outputs)
        {
            foreach(var item in grad_outputs)
            {
                item.Dispose();
            }
            return [null, null, null, null, null, null, null, null, null, null];
        }
    }
}
